Codecov Report

❌ Patch coverage is 99.10714% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 88.39%. Comparing base (e1742e0) to head (108bc13).
⚠️ Report is 11 commits behind head on main.

@@            Coverage Diff             @@
##             main    #2515      +/-   ##
==========================================
+ Coverage   88.37%   88.39%   +0.02%     
==========================================
  Files         620      622       +2     
  Lines      228395   230170    +1775     
==========================================
+ Hits       201836   203457    +1621     
- Misses      26035    26189     +154     
  Partials      524      524              

The sequencing here looks off. In graceful shutdown the runtime does:

DrainIngress -> ReceiverDrained -> downstream Shutdown.

This change makes fake_data_generator do

DrainIngress -> ReceiverDrained -> wait for Shutdown,

but that Shutdown is not part of the normal post-drain receiver path. For this receiver, once ingress is stopped there is no receiver-local work left to preserve, so it should exit directly on DrainIngress rather than report drained and then block waiting for another shutdown message.

The correct fix should be:

DrainIngress -> notify_receiver_drained() -> return TerminalState immediately

Something like (not tested):

  Ok(NodeControlMsg::DrainIngress { deadline, .. }) => {
      otel_info!("fake_data_generator.drain_ingress");                                                                                                                                                                               
      effect_handler.notify_receiver_drained().await?;
      return Ok(TerminalState::new(deadline, [self.metrics.snapshot()]));                                                                                                                                                            
  }    

The fix now looks correct. However from CI failures, there looks like a shutdown race in fake_data_generator that is easy to hit on slower runners The test config uses signals_per_second = 1, so the receiver can sleep for close to 1 second between sends, while the shutdown deadline in test_telemetry_registries_cleanup is only 200ms. That means DrainIngress can arrive while the receiver is asleep, the runtime can move into forced shutdown before the receiver handles it, and then notify_receiver_drained().await? can fail with Channel is closed.

One option could be to address this in two places:

• make the rate-limit sleep interruptible, since that looks like the root cause here.

if signals_per_second.is_some() {
    let remaining_time = wait_till - Instant::now();
    if remaining_time.as_secs_f64() > 0.0 {
        tokio::select! {
            biased;

            ctrl_msg = ctrl_msg_recv.recv() => {
                // handle DrainIngress / Shutdown during the rate-limit wait
                // using the same control-message handling as the main loop
            }

            _ = sleep(remaining_time) => {}
        }
    }
}

• make notify_receiver_drained() best-effort on the terminal DrainIngress path, so a late control-plane teardown does not turn shutdown into an error.

Ok(NodeControlMsg::DrainIngress { deadline, .. }) => {
    otel_info!("fake_data_generator.drain_ingress");
    let _ = effect_handler.notify_receiver_drained().await;
    return Ok(TerminalState::new(deadline, [self.metrics.snapshot()]));
}

@sjmsft

[in addition to the sequence described by @lalitb ]

The exception is deadline-forced shutdown. If the drain deadline expires before the receiver reports drained, the runtime sends NodeControlMsg::Shutdown { deadline, reason } to any still-pending receivers.